#pragma once

#include "PointLabCore/IPointCloud.h"

#include "PointLabTools/Prerequisites.h"
#include "PointLabTools/Kdtree/nanoflann.h"

namespace PointlabTools
{
	using namespace nanoflann;

	class POINTLABTOOLS_EXPORT PointCloudAdaptor
	{
	public:

		PointCloudAdaptor(PointLabCore::IPointCloud *cloud)
		{
			m_cloud = cloud;
		}

		~PointCloudAdaptor() {}

		inline size_t kdtree_get_point_count() const { return m_cloud->getPointNum(); }

		// Returns the distance between the vector "p1[0:size-1]" and the data point with index "idx_p2" stored in the class:
		inline PointLabCore::Real kdtree_distance(const PointLabCore::Real *p1, const size_t idx_p2, size_t /*size*/) const
		{
			const PointLabCore::Real d0 = p1[0] - m_cloud->getPointCoord(idx_p2).x;
			const PointLabCore::Real d1 = p1[1] - m_cloud->getPointCoord(idx_p2).y;
			const PointLabCore::Real d2 = p1[2] - m_cloud->getPointCoord(idx_p2).z;
			return d0*d0 + d1*d1 + d2*d2;
		}

		// Returns the dim'th component of the idx'th point in the class:
		// Since this is inlined and the "dim" argument is typically an immediate value, the
		//  "if/else's" are actually solved at compile time.
		inline PointLabCore::Real kdtree_get_pt(const size_t idx, int dim) const
		{
			if (dim == 0) return m_cloud->getPointCoord(idx).x;
			else if (dim == 1) return m_cloud->getPointCoord(idx).y;
			else return m_cloud->getPointCoord(idx).z;
		}

		// Optional bounding-box computation: return false to default to a standard bbox computation loop.
		//   Return true if the BBOX was already computed by the class and returned in "bb" so it can be avoided to redo it again.
		//   Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3 for point clouds)
		template <class BBOX>
		bool kdtree_get_bbox(BBOX& /*bb*/) const { return false; }
	private:
		PointLabCore::IPointCloud* m_cloud;
	};

	typedef KDTreeSingleIndexAdaptor<
		L2_Simple_Adaptor<PointLabCore::Real, PointCloudAdaptor >,
		PointCloudAdaptor,
		3 /* dim */
	> FlannKDTree;

	class POINTLABTOOLS_EXPORT KDTree :public FlannKDTree
	{
	public:
		KDTree(PointLabCore::IPointCloud* cloud);
		~KDTree();

		void knnSearch(const PointLabCore::Vector3& point, const PointLabCore::Uint k, std::vector<PointLabCore::Uint>& indices, std::vector<PointLabCore::Real>& distance);

		PointLabCore::Uint radiusSearch(const PointLabCore::Vector3& point, const PointLabCore::Real &radius, std::vector<PointLabCore::Uint>& indices, std::vector<PointLabCore::Real>& distance, bool sort);
	};
}
